package org.tim.websocket.protocol;

import lombok.var;
import org.tim.common.ProtocolEnum;
import org.tim.common.TimConstants;
import org.tim.spi.Protocol;
import org.tio.core.ChannelContext;
import org.tio.core.exception.TioDecodeException;
import org.tio.http.common.HttpRequest;
import org.tio.http.common.HttpRequestDecoder;
import org.tio.websocket.server.WsServerAioHandler;
import org.tio.websocket.server.WsServerConfig;
import org.tio.websocket.server.handler.IWsMsgHandler;

import java.nio.ByteBuffer;

public class WebSocketProtocol extends WsServerAioHandler implements Protocol {
    private final WsServerConfig wsServerConfig;
    /**
     * @param wsServerConfig
     * @param wsMsgHandler
     */
    public WebSocketProtocol(WsServerConfig wsServerConfig, IWsMsgHandler wsMsgHandler) {
        super(wsServerConfig, wsMsgHandler);
        this.wsServerConfig = wsServerConfig;
    }

    @Override
    public boolean support(ByteBuffer byteBuffer, int i, int i1, int i2, ChannelContext channelContext) {
        var protocol = channelContext.getAttribute(TimConstants.CONTEXT_PROTOCOL_KEY);
        if (ProtocolEnum.WEBSOCKET.equals(protocol) || isWsProtocol(byteBuffer, i, i1, i2, channelContext)) {
            channelContext.setAttribute(TimConstants.CONTEXT_PROTOCOL_KEY, ProtocolEnum.WEBSOCKET);
            return true;
        }
        return false;
    }

    @Override
    public boolean support(ChannelContext channelContext) {
        var protocol = channelContext.getAttribute(TimConstants.CONTEXT_PROTOCOL_KEY);
        if(protocol != null && ProtocolEnum.WEBSOCKET.equals(protocol)) {
            return true;
        }
        return false;
    }

    private boolean isWsProtocol(ByteBuffer byteBuffer, int i, int i1, int i2, ChannelContext channelContext) {
        HttpRequest request = null;
        try {
            request = HttpRequestDecoder.decode(byteBuffer, i, i1, i2, channelContext, wsServerConfig);

            byteBuffer.rewind();
        } catch (TioDecodeException e) {
            return false;
        }

        if (request != null && request.getHeaders().get("Sec-WebSocket-Key".toLowerCase()) != null) {
            return true;
        }
        return false;
    }
}
